{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "254bbbf2",
   "metadata": {},
   "source": [
    "# Control Flows for JIT compilation\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/brainpy/brainpy/blob/master/docs_version2/tutorial_math/control_flows.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/brainpy/brainpy/blob/master/docs_version2/tutorial_math/control_flows.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "355bb9b6",
   "metadata": {},
   "source": [
    "@[Chaoming Wang](https://github.com/chaoming0625)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8d7be18d8e56574",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Control flow is the core of a program, because it defines the order in which the program's code executes. The control flow of Python is regulated by *conditional statements*, *loops*, and *function calls*.\n",
    "\n",
    "\n",
    "Python has two types of control structures:\n",
    "\n",
    "- **Selection**: used for decisions and branching.\n",
    "- **Repetition**: used for looping, i.e., repeating a piece of code multiple times.\n",
    "\n",
    "\n",
    "In this section, we are going to talk about how to build effective control flows in the context of JIT compilation."
   ]
  },
  {
   "cell_type": "code",
   "id": "38a2bb50",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:29.892608Z",
     "start_time": "2025-10-06T04:00:26.024279Z"
    }
   },
   "source": [
    "import brainpy as bp\n",
    "import brainpy.math as bm\n",
    "\n",
    "bm.set_platform('cpu')"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "cell_type": "code",
   "id": "4dd6ae3c3ac8fb4d",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:29.907013Z",
     "start_time": "2025-10-06T04:00:29.899816Z"
    }
   },
   "source": [
    "bp.__version__"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'3.0.0'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "id": "209cae8bf31357fa",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## 1\\. Selection"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76300eb64d376d38",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "In Python, the selection statements are also known as *Decision control statements* or *branching statements*. The selection statement allows a program to test several conditions and execute instructions based on which condition is true. The commonly used control statements include:\n",
    "\n",
    "- if-else\n",
    "- nested if\n",
    "- if-elif-else"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "143cc7440a437281",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Non-`Variable`-based control statements"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79e9f1a1f7079f82",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Actually, BrainPy (based on JAX) **allows to write control flows normally like your familiar Python programs, when the conditional statement depends on non-Variable instances**. For example,"
   ]
  },
  {
   "cell_type": "code",
   "id": "23fbea620f3eefa2",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:29.930155Z",
     "start_time": "2025-10-06T04:00:29.925651Z"
    }
   },
   "source": [
    "class OddEven(bp.BrainPyObject):\n",
    "    def __init__(self, type_=1):\n",
    "        super(OddEven, self).__init__()\n",
    "        self.type_ = type_\n",
    "        self.a = bm.Variable(bm.zeros(1))\n",
    "\n",
    "    def __call__(self):\n",
    "        if self.type_ == 1:\n",
    "            self.a += 1\n",
    "        elif self.type_ == 2:\n",
    "            self.a -= 1\n",
    "        else:\n",
    "            raise ValueError(f'Unknown type: {self.type_}')\n",
    "        return self.a"
   ],
   "outputs": [],
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "id": "68774c05f83ea6c2",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "In the above example, the target *statement* in ``if (statement)`` syntax relies on a scalar, which is not an instance of [brainpy.math.Variable](./arrays_and_variables.ipynb). In this case, the conditional statements can be arbitrarily complex. You can write your models with normal Python codes. These models will work very well with [JIT compilation](./jit_compilation.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "id": "5e17f40b11e06350",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:29.999383Z",
     "start_time": "2025-10-06T04:00:29.957254Z"
    }
   },
   "source": [
    "model = bm.jit(OddEven(type_=1))\n",
    "\n",
    "model()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1.], dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 4
  },
  {
   "cell_type": "code",
   "id": "1334846bd3ed43",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:30.031850Z",
     "start_time": "2025-10-06T04:00:30.005727Z"
    }
   },
   "source": [
    "model = bm.jit(OddEven(type_=2))\n",
    "\n",
    "model()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([-1.], dtype=float32)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 5
  },
  {
   "cell_type": "code",
   "id": "9dc5b569a2e95727",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:30.060222Z",
     "start_time": "2025-10-06T04:00:30.038805Z"
    }
   },
   "source": [
    "try:\n",
    "    model = bm.jit(OddEven(type_=3))\n",
    "    model()\n",
    "except ValueError as e:\n",
    "    print(f\"ValueError: {str(e)}\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ValueError: Unknown type: 3\n"
     ]
    }
   ],
   "execution_count": 6
  },
  {
   "cell_type": "markdown",
   "id": "c8fd9659fac2317d",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### `Variable`-based control statements"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bb5f79b1c71cf0",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "However, if the `statement` target in a ``if ... else ...`` syntax relies on instances of [brainpy.math.Variable](./arrays_and_variables.ipynb), writing Pythonic control flows will cause errors when using JIT compilation."
   ]
  },
  {
   "cell_type": "code",
   "id": "639811afe3af4e79",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:30.070139Z",
     "start_time": "2025-10-06T04:00:30.066414Z"
    }
   },
   "source": [
    "class OddEvenCauseError(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(OddEvenCauseError, self).__init__()\n",
    "        self.rand = bm.Variable(bm.random.random(1))\n",
    "        self.a = bm.Variable(bm.zeros(1))\n",
    "\n",
    "    def __call__(self):\n",
    "        if self.rand < 0.5:  self.a += 1\n",
    "        else:  self.a -= 1\n",
    "        return self.a"
   ],
   "outputs": [],
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "id": "eac384ce730144be",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:30.258572Z",
     "start_time": "2025-10-06T04:00:30.088899Z"
    }
   },
   "source": [
    "wrong_model = bm.jit(OddEvenCauseError())\n",
    "\n",
    "try:\n",
    "    wrong_model()\n",
    "except Exception as e:\n",
    "    print(f\"{e.__class__.__name__}: {str(e)}\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[1].\n",
      "See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError\n"
     ]
    }
   ],
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "id": "6b79ffb013902560",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "To perform conditional statement on [Variable](./arrays_and_variables.ipynb) instances, we need structural control flow syntax. Specifically, BrainPy provides several options (based on JAX):\n",
    "\n",
    "- [brainpy.math.where](https://numpy.org/doc/stable/reference/generated/numpy.where.html): return element-wise conditional comparison results.\n",
    "- [brainpy.math.ifelse](../apis/auto/generated/brainpy.math.ifelse.rst): Conditional statements of `if-else`, or `if-elif-else`, ... for a scalar-typed value."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "396b4a871994b39d",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### `brainpy.math.where`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8419c2e1b83ef9da",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "``where(condition, x, y)`` function returns elements chosen from *x* or *y* depending on *condition*. It can perform well on scalars, vectors, and high-dimensional arrays."
   ]
  },
  {
   "cell_type": "code",
   "id": "182ae1cae035ca96",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:30.285530Z",
     "start_time": "2025-10-06T04:00:30.263648Z"
    }
   },
   "source": [
    "a = 1.\n",
    "bm.where(a < 0, 0., 1.)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(1., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 9
  },
  {
   "cell_type": "code",
   "id": "24b8510d06a06812",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:30.432160Z",
     "start_time": "2025-10-06T04:00:30.304491Z"
    }
   },
   "source": [
    "a = bm.random.random(5)\n",
    "bm.where(a < 0.5, 0., 1.)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([1., 1., 0., 0., 1.]), dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 10
  },
  {
   "cell_type": "code",
   "id": "1f6d7e039fb7421c",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:30.583838Z",
     "start_time": "2025-10-06T04:00:30.452768Z"
    }
   },
   "source": [
    "a = bm.random.random((3, 3))\n",
    "bm.where(a < 0.5, 0., 1.)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([[0., 1., 0.],\n",
       "                   [1., 0., 1.],\n",
       "                   [1., 1., 1.]]),\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 11
  },
  {
   "cell_type": "markdown",
   "id": "15a3152c3e6c8117",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "For the above example, we can rewrite it by using `where` syntax as:"
   ]
  },
  {
   "cell_type": "code",
   "id": "be48e4fc29e6cfbe",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:30.592396Z",
     "start_time": "2025-10-06T04:00:30.589025Z"
    }
   },
   "source": [
    "class OddEvenWhere(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(OddEvenWhere, self).__init__()\n",
    "        self.rand = bm.Variable(bm.random.random(1))\n",
    "        self.a = bm.Variable(bm.zeros(1))\n",
    "\n",
    "    def __call__(self):\n",
    "        self.a += bm.where(self.rand < 0.5, 1., -1.)\n",
    "        return self.a"
   ],
   "outputs": [],
   "execution_count": 12
  },
  {
   "cell_type": "code",
   "id": "99a8a56dd22ff7b8",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:30.644021Z",
     "start_time": "2025-10-06T04:00:30.608950Z"
    }
   },
   "source": [
    "model = bm.jit(OddEvenWhere())\n",
    "model()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([-1.], dtype=float32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 13
  },
  {
   "cell_type": "markdown",
   "id": "53a50cafa99b85cc",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### `brainpy.math.ifelse`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f67dbf5043879cdf",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Based on JAX's control flow syntax [jax.lax.cond](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.cond.html), BrainPy provides a more general conditional statement enabling multiple branching.\n",
    "\n",
    "In its simplest case, `brainpy.math.ifelse(condition, branches, operands)` is equivalent to:\n",
    "\n",
    "```python\n",
    "def ifelse(condition, branches, operands):\n",
    "  true_fun, false_fun = branches\n",
    "  if condition:\n",
    "    return true_fun(operands)\n",
    "  else:\n",
    "    return false_fun(operands)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96504b2a8e5cb3f9",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Based on this function, we can rewrite the above example by using `cond` syntax as:"
   ]
  },
  {
   "cell_type": "code",
   "id": "87052073a05de81b",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:30.653504Z",
     "start_time": "2025-10-06T04:00:30.646897Z"
    }
   },
   "source": [
    "class OddEvenCond(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(OddEvenCond, self).__init__()\n",
    "        self.rand = bm.Variable(bm.random.random(1))\n",
    "        self.a = bm.Variable(bm.zeros(1))\n",
    "\n",
    "    def __call__(self):\n",
    "        self.a += bm.ifelse(self.rand[0] < 0.5,\n",
    "                            [lambda: 1., lambda: -1.])\n",
    "        return self.a"
   ],
   "outputs": [],
   "execution_count": 14
  },
  {
   "cell_type": "code",
   "id": "b2bc687032e46f6e",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.238378Z",
     "start_time": "2025-10-06T04:00:30.670055Z"
    }
   },
   "source": [
    "model = bm.jit(OddEvenCond())\n",
    "model()"
   ],
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "len() of unsized object",
     "output_type": "error",
     "traceback": [
      "\u001B[31m---------------------------------------------------------------------------\u001B[39m",
      "\u001B[31mIndexError\u001B[39m                                Traceback (most recent call last)",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\jax\\_src\\core.py:2234\u001B[39m, in \u001B[36mShapedArray._len\u001B[39m\u001B[34m(self, ignored_tracer)\u001B[39m\n\u001B[32m   2233\u001B[39m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[32m-> \u001B[39m\u001B[32m2234\u001B[39m   \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m.shape[\u001B[32m0\u001B[39m]\n\u001B[32m   2235\u001B[39m \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mIndexError\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m err:\n",
      "\u001B[31mIndexError\u001B[39m: tuple index out of range",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001B[31mTypeError\u001B[39m                                 Traceback (most recent call last)",
      "\u001B[36mCell\u001B[39m\u001B[36m \u001B[39m\u001B[32mIn[15]\u001B[39m\u001B[32m, line 2\u001B[39m\n\u001B[32m      1\u001B[39m model = bm.jit(OddEvenCond())\n\u001B[32m----> \u001B[39m\u001B[32m2\u001B[39m model()\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_jit.py:97\u001B[39m, in \u001B[36m_get_jitted_fun.<locals>.jitted_fun\u001B[39m\u001B[34m(*args, **params)\u001B[39m\n\u001B[32m     94\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m fun.fun(*args, **params)\n\u001B[32m     96\u001B[39m \u001B[38;5;66;03m# compile the function and get the state trace\u001B[39;00m\n\u001B[32m---> \u001B[39m\u001B[32m97\u001B[39m state_trace = fun.get_state_trace_by_call(*args, **params)\n\u001B[32m     98\u001B[39m read_state_vals = state_trace.get_read_state_values(\u001B[38;5;28;01mTrue\u001B[39;00m)\n\u001B[32m    100\u001B[39m \u001B[38;5;66;03m# call the jitted function\u001B[39;00m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:689\u001B[39m, in \u001B[36mStatefulFunction.get_state_trace_by_call\u001B[39m\u001B[34m(self, compile_if_miss, *args, **kwargs)\u001B[39m\n\u001B[32m    671\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34mget_state_trace_by_call\u001B[39m(\u001B[38;5;28mself\u001B[39m, *args, compile_if_miss: \u001B[38;5;28mbool\u001B[39m = \u001B[38;5;28;01mTrue\u001B[39;00m, **kwargs) -> StateTraceStack:\n\u001B[32m    672\u001B[39m \u001B[38;5;250m    \u001B[39m\u001B[33;03m\"\"\"\u001B[39;00m\n\u001B[32m    673\u001B[39m \u001B[33;03m    Read the state trace of the function.\u001B[39;00m\n\u001B[32m    674\u001B[39m \n\u001B[32m   (...)\u001B[39m\u001B[32m    687\u001B[39m \u001B[33;03m        The state trace of the function.\u001B[39;00m\n\u001B[32m    688\u001B[39m \u001B[33;03m    \"\"\"\u001B[39;00m\n\u001B[32m--> \u001B[39m\u001B[32m689\u001B[39m     cache_key = \u001B[38;5;28mself\u001B[39m.get_arg_cache_key(*args, **kwargs, compile_if_miss=compile_if_miss)\n\u001B[32m    690\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m.get_state_trace(cache_key)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:839\u001B[39m, in \u001B[36mStatefulFunction.get_arg_cache_key\u001B[39m\u001B[34m(self, compile_if_miss, *args, **kwargs)\u001B[39m\n\u001B[32m    831\u001B[39m cache_key = hashabledict(\n\u001B[32m    832\u001B[39m     static_args=static_args,\n\u001B[32m    833\u001B[39m     dyn_args=dyn_args,\n\u001B[32m    834\u001B[39m     static_kwargs=static_kwargs,\n\u001B[32m    835\u001B[39m     dyn_kwargs=dyn_kwargs,\n\u001B[32m    836\u001B[39m )\n\u001B[32m    838\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m cache_key \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mself\u001B[39m._cached_state_trace \u001B[38;5;129;01mand\u001B[39;00m compile_if_miss:\n\u001B[32m--> \u001B[39m\u001B[32m839\u001B[39m     \u001B[38;5;28mself\u001B[39m.make_jaxpr(*args, **kwargs)\n\u001B[32m    841\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m cache_key\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:950\u001B[39m, in \u001B[36mStatefulFunction.make_jaxpr\u001B[39m\u001B[34m(self, *args, **kwargs)\u001B[39m\n\u001B[32m    948\u001B[39m         \u001B[38;5;28mself\u001B[39m._cached_out_shapes.pop(cache_key, \u001B[38;5;28;01mNone\u001B[39;00m)\n\u001B[32m    949\u001B[39m         \u001B[38;5;28mself\u001B[39m._cached_jaxpr.pop(cache_key, \u001B[38;5;28;01mNone\u001B[39;00m)\n\u001B[32m--> \u001B[39m\u001B[32m950\u001B[39m         \u001B[38;5;28;01mraise\u001B[39;00m e\n\u001B[32m    952\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:928\u001B[39m, in \u001B[36mStatefulFunction.make_jaxpr\u001B[39m\u001B[34m(self, *args, **kwargs)\u001B[39m\n\u001B[32m    926\u001B[39m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[32m    927\u001B[39m         dyn_kwargs[k] = v\n\u001B[32m--> \u001B[39m\u001B[32m928\u001B[39m jaxpr, (out_shapes, state_shapes) = _make_jaxpr(\n\u001B[32m    929\u001B[39m     functools.partial(\n\u001B[32m    930\u001B[39m         \u001B[38;5;28mself\u001B[39m._wrapped_fun_to_eval,\n\u001B[32m    931\u001B[39m         cache_key,\n\u001B[32m    932\u001B[39m         static_kwargs,\n\u001B[32m    933\u001B[39m         return_only_write=\u001B[38;5;28mself\u001B[39m.return_only_write,\n\u001B[32m    934\u001B[39m     ),\n\u001B[32m    935\u001B[39m     static_argnums=\u001B[38;5;28mself\u001B[39m.static_argnums,\n\u001B[32m    936\u001B[39m     axis_env=\u001B[38;5;28mself\u001B[39m.axis_env,\n\u001B[32m    937\u001B[39m     return_shape=\u001B[38;5;28;01mTrue\u001B[39;00m,\n\u001B[32m    938\u001B[39m     abstracted_axes=\u001B[38;5;28mself\u001B[39m.abstracted_axes,\n\u001B[32m    939\u001B[39m )(*args, **dyn_kwargs)\n\u001B[32m    940\u001B[39m \u001B[38;5;66;03m# returns\u001B[39;00m\n\u001B[32m    941\u001B[39m \u001B[38;5;28mself\u001B[39m._cached_jaxpr_out_tree.set(cache_key, jax.tree.structure((out_shapes, state_shapes)))\n",
      "    \u001B[31m[... skipping hidden 1 frame]\u001B[39m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:1399\u001B[39m, in \u001B[36m_make_jaxpr.<locals>.make_jaxpr_f\u001B[39m\u001B[34m(*args, **kwargs)\u001B[39m\n\u001B[32m   1397\u001B[39m         jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info_)\n\u001B[32m   1398\u001B[39m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[32m-> \u001B[39m\u001B[32m1399\u001B[39m         jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f)\n\u001B[32m   1400\u001B[39m closed_jaxpr = ClosedJaxpr(jaxpr, consts)\n\u001B[32m   1401\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m return_shape:\n",
      "    \u001B[31m[... skipping hidden 5 frame]\u001B[39m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:874\u001B[39m, in \u001B[36mStatefulFunction._wrapped_fun_to_eval\u001B[39m\u001B[34m(self, cache_key, static_kwargs, return_only_write, *args, **dyn_kwargs)\u001B[39m\n\u001B[32m    872\u001B[39m \u001B[38;5;28mself\u001B[39m._cached_state_trace.set(cache_key, state_trace)\n\u001B[32m    873\u001B[39m \u001B[38;5;28;01mwith\u001B[39;00m state_trace:\n\u001B[32m--> \u001B[39m\u001B[32m874\u001B[39m     out = \u001B[38;5;28mself\u001B[39m.fun(*args, **dyn_kwargs, **static_kwargs)\n\u001B[32m    875\u001B[39m     state_values = (\n\u001B[32m    876\u001B[39m         state_trace.get_write_state_values(\u001B[38;5;28;01mTrue\u001B[39;00m)\n\u001B[32m    877\u001B[39m         \u001B[38;5;28;01mif\u001B[39;00m return_only_write \u001B[38;5;28;01melse\u001B[39;00m\n\u001B[32m    878\u001B[39m         state_trace.get_state_values()\n\u001B[32m    879\u001B[39m     )\n\u001B[32m    880\u001B[39m state_trace.recovery_original_values()\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\BrainPy\\brainpy\\version2\\math\\object_transform\\_utils.py:68\u001B[39m, in \u001B[36mwarp_to_no_state_input_output.<locals>.wrapper\u001B[39m\u001B[34m(*args, **kwargs)\u001B[39m\n\u001B[32m     65\u001B[39m \u001B[38;5;129m@wraps\u001B[39m(fn)\n\u001B[32m     66\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34mwrapper\u001B[39m(*args, **kwargs):\n\u001B[32m     67\u001B[39m     args, kwargs = jax.tree.map(_remove_state, (args, kwargs), is_leaf=\u001B[38;5;28;01mlambda\u001B[39;00m x: \u001B[38;5;28misinstance\u001B[39m(x, brainstate.State))\n\u001B[32m---> \u001B[39m\u001B[32m68\u001B[39m     out = fn(*args, **kwargs)\n\u001B[32m     69\u001B[39m     out = jax.tree.map(_remove_state, out, is_leaf=\u001B[38;5;28;01mlambda\u001B[39;00m x: \u001B[38;5;28misinstance\u001B[39m(x, brainstate.State))\n\u001B[32m     70\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m out\n",
      "\u001B[36mCell\u001B[39m\u001B[36m \u001B[39m\u001B[32mIn[14]\u001B[39m\u001B[32m, line 8\u001B[39m, in \u001B[36mOddEvenCond.__call__\u001B[39m\u001B[34m(self)\u001B[39m\n\u001B[32m      7\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34m__call__\u001B[39m(\u001B[38;5;28mself\u001B[39m):\n\u001B[32m----> \u001B[39m\u001B[32m8\u001B[39m     \u001B[38;5;28mself\u001B[39m.a += bm.ifelse(\u001B[38;5;28mself\u001B[39m.rand[\u001B[32m0\u001B[39m] < \u001B[32m0.5\u001B[39m,\n\u001B[32m      9\u001B[39m                         [\u001B[38;5;28;01mlambda\u001B[39;00m: \u001B[32m1.\u001B[39m, \u001B[38;5;28;01mlambda\u001B[39;00m: -\u001B[32m1.\u001B[39m])\n\u001B[32m     10\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m.a\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\BrainPy\\brainpy\\version2\\math\\object_transform\\controls.py:187\u001B[39m, in \u001B[36mifelse\u001B[39m\u001B[34m(conditions, branches, operands)\u001B[39m\n\u001B[32m    183\u001B[39m         exclusive_conditions.append(all_false)\n\u001B[32m    185\u001B[39m     conditions = exclusive_conditions\n\u001B[32m--> \u001B[39m\u001B[32m187\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m brainstate.transform.ifelse(conditions, branches, *operands)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_conditions.py:306\u001B[39m, in \u001B[36mifelse\u001B[39m\u001B[34m(conditions, branches, check_cond, *operands)\u001B[39m\n\u001B[32m    304\u001B[39m \u001B[38;5;28;01melif\u001B[39;00m \u001B[38;5;28mlen\u001B[39m(branches) == \u001B[32m1\u001B[39m:\n\u001B[32m    305\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m branches[\u001B[32m0\u001B[39m](*operands)\n\u001B[32m--> \u001B[39m\u001B[32m306\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mlen\u001B[39m(conditions) != \u001B[38;5;28mlen\u001B[39m(branches):\n\u001B[32m    307\u001B[39m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mValueError\u001B[39;00m(\u001B[33m\"\u001B[39m\u001B[33mThe number of conditions should be equal to the number of branches.\u001B[39m\u001B[33m\"\u001B[39m)\n\u001B[32m    309\u001B[39m \u001B[38;5;66;03m# format index\u001B[39;00m\n",
      "    \u001B[31m[... skipping hidden 1 frame]\u001B[39m\n",
      "\u001B[36mFile \u001B[39m\u001B[32m~\\miniconda3\\envs\\bdp\\Lib\\site-packages\\jax\\_src\\core.py:2236\u001B[39m, in \u001B[36mShapedArray._len\u001B[39m\u001B[34m(self, ignored_tracer)\u001B[39m\n\u001B[32m   2234\u001B[39m   \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m.shape[\u001B[32m0\u001B[39m]\n\u001B[32m   2235\u001B[39m \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mIndexError\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m err:\n\u001B[32m-> \u001B[39m\u001B[32m2236\u001B[39m   \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mTypeError\u001B[39;00m(\u001B[33m\"\u001B[39m\u001B[33mlen() of unsized object\u001B[39m\u001B[33m\"\u001B[39m) \u001B[38;5;28;01mfrom\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34;01merr\u001B[39;00m\n",
      "\u001B[31mTypeError\u001B[39m: len() of unsized object"
     ]
    }
   ],
   "execution_count": 15
  },
  {
   "cell_type": "markdown",
   "id": "4a755a6feaea6da1",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "If you want to write control flows with multiple branchings, ``brainpy.math.ifelse(conditions, branches, operands)`` can also help you accomplish this easily. Actually, multiple branching case is equivalent to:\n",
    "\n",
    "```python\n",
    "def ifelse(conditions, branches, operands):\n",
    "  pred1, pred2, ... = conditions\n",
    "  func1, func2, ..., funcN = branches\n",
    "  if pred1:\n",
    "    return func1(operands)\n",
    "  elif pred2:\n",
    "    return func2(operands)\n",
    "  ...\n",
    "  else:\n",
    "    return funcN(operands)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76268713ca8c307b",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "For example, if you have the following code:\n",
    "\n",
    "```python\n",
    "def f(a):\n",
    "  if a > 10:\n",
    "    return 1.\n",
    "  elif a > 5:\n",
    "    return 2.\n",
    "  elif a > 0:\n",
    "    return 3.\n",
    "  elif a > -5:\n",
    "    return 4.\n",
    "  else:\n",
    "    return 5.\n",
    "```\n",
    "\n",
    "It can be expressed as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "904d2ef57be4cc6a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.240385800Z",
     "start_time": "2023-11-05T01:49:50.767081400Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def f(a):\n",
    "  return bm.ifelse(conditions=[a > 10, a > 5, a > 0, a > -5],\n",
    "                   branches=[1., 2., 3., 4., 5.])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "22882e5c3c7b91b3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.241393300Z",
     "start_time": "2023-11-05T01:49:50.773526Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(1., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(11.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2f8b2eae8a8e7666",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.241393300Z",
     "start_time": "2023-11-05T01:49:50.819594200Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(2., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(6.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "fbad099531113d56",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.241393300Z",
     "start_time": "2023-11-05T01:49:50.838919Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(3., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "79be4f3c51590366",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.241393300Z",
     "start_time": "2023-11-05T01:49:50.864183500Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(4., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(-4.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "1c6eb117f19d6a62",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.241393300Z",
     "start_time": "2023-11-05T01:49:50.895823700Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(5., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(-6.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfd50ae5cb48a2a5",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "A more complex example is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ed202700dd2bf611",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.242801500Z",
     "start_time": "2023-11-05T01:49:50.921011400Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def f2(a, x):\n",
    "  return bm.ifelse(conditions=[a > 10, a > 5, a > 0, a > -5],\n",
    "                   branches=[lambda x: x*2,\n",
    "                             2.,\n",
    "                             lambda x: x**2 -1,\n",
    "                             lambda x: x - 4.,\n",
    "                             5.],\n",
    "                   operands=x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b1cae24c7aa5e5d8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.242801500Z",
     "start_time": "2023-11-05T01:49:50.922795400Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(2., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f2(11, 1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c1fd9c9bad31be6d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.243811400Z",
     "start_time": "2023-11-05T01:49:50.958944Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(2., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f2(6, 1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "2a5ca3a331aaedda",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.247808500Z",
     "start_time": "2023-11-05T01:49:50.996158400Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f2(1, 1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "680e8e688b234181",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.247808500Z",
     "start_time": "2023-11-05T01:49:51.021265500Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-3., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f2(-4, 1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "7b0dcbd32262035d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.248811400Z",
     "start_time": "2023-11-05T01:49:51.056306100Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(5., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f2(-6, 1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a4fefd5188501d6",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "If instances of `brainpy.math.Variable` are used in branching functions,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "4618ac0a007d43",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.248811400Z",
     "start_time": "2023-11-05T01:49:51.086629Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a: Variable(value=Array([1., 1.]), dtype=float32)\n",
      "b: Variable(value=Array([0., 0.]), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "a = bm.Variable(bm.zeros(2))\n",
    "b = bm.Variable(bm.ones(2))\n",
    "def true_f():  a.value += 1\n",
    "def false_f(): b.value -= 1\n",
    "\n",
    "bm.ifelse(True, [true_f, false_f])\n",
    "bm.ifelse(False, [true_f, false_f])\n",
    "\n",
    "print('a:', a)\n",
    "print('b:', b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1359741dcab11af4",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## 2\\. Repetition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d61f81207f43640",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "A repetition statement is used to repeat a group(block) of programming instructions.\n",
    "\n",
    "In Python, we generally have two loops/repetitive statements:\n",
    "\n",
    "- **for loop**: Execute a set of statements once for each item in a sequence.\n",
    "- **while loop**: Execute a block of statements repeatedly until a given condition is satisfied."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "886b0569ed325524",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Pythonic loop syntax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ea62439c2fc0d7a",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Actually, JAX enables to write Pythonic loops. You just need to iterate over you sequence data and then apply your logic on the iterated items. Such kind of Pythonic loop syntax can be compatible with JIT compilation, but will cause long time to trace and compile. For example,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "aaf0042103cfad8f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.259812100Z",
     "start_time": "2023-11-05T01:49:51.154597600Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class LoopSimple(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(LoopSimple, self).__init__()\n",
    "        rng = bm.random.RandomState(123)\n",
    "        self.seq = bm.Variable(rng.random(1000))\n",
    "        self.res = bm.Variable(bm.zeros(1))\n",
    "\n",
    "    def __call__(self):\n",
    "        for s in self.seq:\n",
    "            self.res += s\n",
    "        return self.res.value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "a36cd9996bd0baeb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.276869300Z",
     "start_time": "2023-11-05T01:49:51.156115200Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def measure_time(f, return_res=False, verbose=True):\n",
    "    t0 = time.time()\n",
    "    r = f()\n",
    "    t1 = time.time()\n",
    "    if verbose:\n",
    "        print(f'Result: {r}, Time: {t1 - t0}')\n",
    "    return r if return_res else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "ee410af7c23fb304",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.277867Z",
     "start_time": "2023-11-05T01:49:51.160358900Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Result: [501.74664], Time: 1.4419348239898682\n"
     ]
    }
   ],
   "source": [
    "model = bm.jit(LoopSimple())\n",
    "\n",
    "# First time will trigger compilation\n",
    "measure_time(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "15bf203f8e06331",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.278865Z",
     "start_time": "2023-11-05T01:49:52.750845600Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Result: [1003.4931], Time: 0.0\n"
     ]
    }
   ],
   "source": [
    "# Second running\n",
    "measure_time(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27fe11f02c8098c8",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "When the model is complex and the iteration is long, the compilation during the first running will become unbearable. For such cases, you need structural loop syntax."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13710ef9efd6582c",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "JAX has provided several important loop syntax, including:\n",
    "- [jax.lax.fori_loop](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.fori_loop.html)\n",
    "- [jax.lax.scan](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html)\n",
    "- [jax.lax.while_loop](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.while_loop.html)\n",
    "\n",
    "BrainPy also provides its own loop syntax, which is especially suitable for the cases where users are using `brainpy.math.Variable`. Specifically, they are:\n",
    "\n",
    "- [brainpy.math.for_loop](../apis/auto/math/generated/brainpy.math.object_transform.for_loop.html)\n",
    "- [brainpy.math.while_loop](../apis/auto/math/generated/brainpy.math.object_transform.while_loop.html)\n",
    "\n",
    "In this section, we only talk about how to use our provided loop functions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d99e3e68b7da44b",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### ``brainpy.math.for_loop()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46eb5c901e88c5c0",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "``brainpy.math.for_loop()`` is used to generate a for-loop function when you use ``Variable``.\n",
    "\n",
    "Suppose that you are using several `Variables` to implement your body function \"body\\_fun\". Sometimes the body function already returns something, and you also want to gather the returned values. With the Python syntax, it can be realized as\n",
    "\n",
    "```python\n",
    "\n",
    "def for_loop_function(body_fun, xs):\n",
    "  ys = []\n",
    "  for x in xs:\n",
    "    results = body_fun(x)\n",
    "    ys.append(results)\n",
    "  return ys\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecbfb0219080fe2c",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "In BrainPy, you can define this logic using ``brainpy.math.for_loop()``:\n",
    "\n",
    "```python\n",
    "import brainpy.math\n",
    "\n",
    "hist_of_out_vars = brainpy.math.for_loop(body_fun, operands)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b55243150e40232",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "For the above example, we can rewrite it by using ``brainpy.math.for_loop`` as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "5a40dbcbd5dda52a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.278865Z",
     "start_time": "2023-11-05T01:49:52.760981600Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class LoopStruct(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(LoopStruct, self).__init__()\n",
    "        rng = bm.random.RandomState(123)\n",
    "        self.seq = rng.random(1000)\n",
    "        self.res = bm.Variable(bm.zeros(1))\n",
    "\n",
    "    def __call__(self):\n",
    "        def add(s):\n",
    "          self.res += s\n",
    "          return self.res.value\n",
    "\n",
    "        return bm.for_loop(body_fun=add, operands=self.seq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "9dc917eb00d0f355",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.278865Z",
     "start_time": "2023-11-05T01:49:52.787409300Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1000, 1)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = bm.jit(LoopStruct())\n",
    "\n",
    "r = measure_time(model, verbose=False, return_res=True)\n",
    "r.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "162e91cc18307775",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "In essence, ``body_fun`` defines the one-step updating rule of how variables are updated. All returns of ``body_fun`` will be gathered as the history values. ``operands`` specified the inputs of the ``body_fun``. It will be looped over the fist axis."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a7ad1419077add",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### ``brainpy.math.while_loop()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dac737984aec13f7",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "``brainpy.math.while_loop()`` is used to generate a while-loop function when you use ``Varible``. It supports the following loop logic:\n",
    "\n",
    "```python\n",
    "\n",
    "while condition:\n",
    "    statements\n",
    "```\n",
    "\n",
    "When using ``brainpy.math.while_loop()`` , *condition* should be wrapped as a ``cond_fun`` function which returns a boolean value, and *statements* should be packed as a ``body_fun`` function which receives the old values at the latest step and returns the updated values at the current step:\n",
    "\n",
    "```python\n",
    "\n",
    "while cond_fun(x):\n",
    "    x = body_fun(x)\n",
    "```\n",
    "\n",
    "Note the difference between ``brainpy.math.for_loop`` and ``brainpy.math.while_loop``:\n",
    "\n",
    "1. The returns of ``brainpy.math.for_loop`` are the values to be gathered as the history values. While the returns of ``brainpy.math.while_loop`` should be the same shape and type with its inputs, because they are represented as the updated values.\n",
    "2. ``brainpy.math.for_loop`` can receive anything without explicit requirements of returns. But, ``brainpy.math.while_loop`` should return what it receives.\n",
    "\n",
    "A concreate example of ``brainpy.math.while_loop`` is as the follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "561686dc3dad2d38",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.286948700Z",
     "start_time": "2023-11-05T01:49:52.828787600Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Variable(value=Array([55.]), dtype=float32) Variable(value=Array([10.]), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "i = bm.Variable(bm.zeros(1))\n",
    "counter = bm.Variable(bm.zeros(1))\n",
    "\n",
    "def cond_f():\n",
    "    return i[0] < 10\n",
    "\n",
    "def body_f():\n",
    "    i.value += 1.\n",
    "    counter.value += i\n",
    "\n",
    "bm.while_loop(body_f, cond_f, operands=())\n",
    "\n",
    "print(counter, i)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b0853b62ab6b1d1",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "In the above example, we try to implement a sum from 0 to 10 by using two JaxArrays ``i`` and ``counter``."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b82ac5cfb078a4a",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Or, similarly,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "5b9e31e7fc898515",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:00:33.287945500Z",
     "start_time": "2023-11-05T01:49:52.858201900Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array(56., dtype=float32),)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "i = bm.Variable(bm.zeros(1))\n",
    "\n",
    "def cond_f(counter):\n",
    "    return i[0] < 10\n",
    "\n",
    "def body_f(counter):\n",
    "    i.value += 1.\n",
    "    return counter + i[0]\n",
    "\n",
    "bm.while_loop(body_f, cond_f, operands=(1., ))"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "main_language": "python"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "279.273px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
